import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LinearRegression
from sklearn.metrics import mean_squared_error, r2_score
from sklearn.preprocessing import StandardScaler
import joblib
#假设3-8质量一起预测
# 加载数据
data = pd.read_csv('winequality_red.csv')
# 分离特征与目标变量
selected_features = ['alcohol', 'sulphates', 'volatile acidity', 'citric acid']
def model(select,data,i,mse_list,r2_list,model_save_path=None):
    X = data[selected_features]
    y = data['quality']
    scaler = StandardScaler()  # 数据标准化
    X_scaled = scaler.fit_transform(X)
    X_train, X_test, y_train, y_test = train_test_split(X_scaled, y, test_size=0.2)
    # 创建并训练线性回归模型
    model = LinearRegression()
    model.fit(X_train, y_train)
    # 预测
    y_pred = np.around(model.predict(X_test))
    # 计算MSE和R^2分数
    mse_list[i].append(mean_squared_error(y_test, y_pred))
    r2_list[i].append(r2_score(y_test, y_pred))
    if model_save_path:
        joblib.dump(model, model_save_path)
        print(f"模型已保存至: {model_save_path}")
    # 如果你想查看具体预测值与实际值对比
    # results = pd.DataFrame({'Actual': y_test, 'Predicted': y_pred})
    # print(results.head())
num_runs = 30  # 运行次数
mse_list = {0:[], 1:[], 2:[], 3:[]}
r2_list = {0:[], 1:[], 2:[], 3:[]}
# 分开预测，3-4,5-6,7-8
data_low=data[data['quality']<=4]
data_medium = data[(data['quality'] <= 6) & (data['quality'] > 4)]
data_high = data[data['quality'] >= 7]
for i in range(num_runs):#取不同训练集与测试集运算平均值
    if i==num_runs-1:
        model(selected_features, data, 0, mse_list, r2_list,model_save_path='model0.pkl')
        model(selected_features, data_low, 1, mse_list, r2_list,model_save_path='model1.pkl')
        # 5-6
        model(selected_features, data_medium, 2, mse_list, r2_list,model_save_path='model2.pkl')
        # 7-8
        model(selected_features, data_high, 3, mse_list, r2_list,model_save_path='model3.pkl')
    else:
        model(selected_features, data, 0,mse_list,r2_list)
        model(selected_features, data_low, 1,mse_list,r2_list)
        #5-6
        model(selected_features, data_medium, 2,mse_list,r2_list)
        #7-8
        model(selected_features, data_high, 3,mse_list,r2_list)
mse_mean=[]
r2_mean=[]
for i in range(4):
    mse_mean.append( np.mean(mse_list[i]))
    r2_mean.append(np.mean(r2_list[i]))
print(f"3-8的Mean Squared Error: {mse_mean[0]},r2 Score: {r2_mean[0]}")
for i in range(1,4):
    print(f"{1+2*i}-{1+2*i+1}的Mean Squared Error: {mse_mean[i]},r2 Score: {r2_mean[i]}")
a=input('请输入alcohol, sulphates, volatile acidity, citric acid(用逗号连接）').split(',')
b=input('输入你想调用的模型')
float_list = [[float(item) for item in a]]
model = joblib.load(b)
print('预测值是：',model.predict(float_list))